#include "AVFormatContext.h"


CAVFormatContext::CAVFormatContext(const char *pURI):pFmtCtx(NULL),nVideoIndex(-1),nAudioIndex(-1)
{
	strURI = pURI;
}


CAVFormatContext::~CAVFormatContext(void)
{
}

int CAVFormatContext::Init()
{
	if ((avformat_open_input(&this->pFmtCtx, this->strURI.c_str(), 0, 0)) != 0) 
	{
		printf( "Could not open input file.");
		return -1;
	}

	if ((avformat_find_stream_info(this->pFmtCtx, 0)) != 0)
	{
		printf( "Failed to retrieve input stream information");
		return -1;
	}

	for(unsigned int i = 0; i < this->pFmtCtx->nb_streams; i++) 
	{
		if(this->pFmtCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO)
		{
			this->nVideoIndex = i;
		}
		else if(this->pFmtCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_AUDIO)
		{
			this->nAudioIndex = i;
		}
	}

	av_dump_format(this->pFmtCtx, 0, this->strURI.c_str(), 0);

	return 0;
}

int CAVFormatContext::read_frame()
{
	return av_read_frame(this->pFmtCtx, &this->pkt);
}

void CAVFormatContext::free_packet()
{
	av_free_packet(&this->pkt);
}

AVPacket* CAVFormatContext::getPacket()
{
	return &this->pkt;
}

bool CAVFormatContext::IsVideoPacket()
{
	return this->nVideoIndex == this->pkt.stream_index;
}

bool CAVFormatContext::IsAudioPacket()
{
	return this->nAudioIndex == this->pkt.stream_index;
}

int CAVFormatContext::Init_Output(AVFormatContext *pFmtCtxIn)
{
	avformat_alloc_output_context2(&this->pFmtCtx, NULL, NULL, this->strURI.c_str());
	if (!this->pFmtCtx)
	{
		printf( "Could not create output context\n");
		return -1;
	}

	AVOutputFormat *ofmt = this->pFmtCtx->oformat;
	for (unsigned int i = 0; i < pFmtCtxIn->nb_streams; ++i)
	{
		//Create output AVStream according to input AVStream
		AVStream *in_stream = pFmtCtxIn->streams[i];
		AVStream *out_stream = avformat_new_stream(this->pFmtCtx, in_stream->codec->codec);
		if (!out_stream)
		{
			printf( "Failed allocating output stream\n");
			return -1;
		}
		//Copy the settings of AVCodecContext
		if (avcodec_copy_context(out_stream->codec, in_stream->codec) < 0) 
		{
			printf( "Failed to copy context from input to output stream codec context\n");
			return -1;
		}
		out_stream->codec->codec_tag = 0;
		if (this->pFmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
			out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
	}

	//-Open output file
	if (!(ofmt->flags & AVFMT_NOFILE)) 
	{
		int ret = avio_open(&this->pFmtCtx->pb, this->strURI.c_str(), AVIO_FLAG_WRITE);
		if (ret < 0) 
		{
			printf( "Could not open output file '%s'", this->strURI.c_str());
			return -1;
		}
	}
	//------------------Output information------------------
	av_dump_format(this->pFmtCtx, 0, this->strURI.c_str(), 1);

	return 0;
}

int CAVFormatContext::write_header()
{
	return avformat_write_header(this->pFmtCtx, NULL);
}

int CAVFormatContext::write_frame(AVPacket *pkt)
{
	return av_interleaved_write_frame(this->pFmtCtx, pkt);
}

int CAVFormatContext::write_trailer()
{
	return av_write_trailer(this->pFmtCtx);
}

int CAVFormatContext::Add_Stream( AVFormatContext *pFmtCtxIn )
{
	for (unsigned int i = 0; i < pFmtCtxIn->nb_streams; ++i)
	{
		//Create output AVStream according to input AVStream
		AVStream *in_stream = pFmtCtxIn->streams[i];
		AVStream *out_stream = avformat_new_stream(this->pFmtCtx, in_stream->codec->codec);
		if (!out_stream)
		{
			printf( "Failed allocating output stream\n");
			return -1;
		}
		//Copy the settings of AVCodecContext
		if (avcodec_copy_context(out_stream->codec, in_stream->codec) < 0) 
		{
			printf( "Failed to copy context from input to output stream codec context\n");
			return -1;
		}
		out_stream->codec->codec_tag = 0;
		if (this->pFmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
			out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
	}

	//------------------Output information------------------
	av_dump_format(this->pFmtCtx, 0, this->strURI.c_str(), 1);

	return 0;
}
